import os
from PIL import Image
import cv2
import random
import argparse
from torchvision import transforms

# 随机分割
def split_augment_image_random(path, file, output, size=1024, num=12):
    image = Image.open(path)
    image_aug = transforms.Compose([
        transforms.RandomHorizontalFlip(),
        transforms.RandomVerticalFlip(),
        transforms.RandomCrop(size)
    ])
    _class = path.split("/")[-2]
    filename = file.split(".")
    for i in range(num):
        image_aug(image).save("Photos/{}/{}/{}_{}.{}".format(output, _class, filename[0], str(i).zfill(2), filename[1]))

# 滑动分割
def split_augment_image(path, file, output):
    image = cv2.imread(path)
    heigh = image.shape[0]
    width = image.shape[1]
    _class = path.split("/")[-2]
    filename = file.split(".")
    count = 0
    for w in range(0, width-256, 256):
        for h in range(0, heigh-256, 256):
            count += 1
            result_image = image[h:h+512,w:w+512]

            # [0, 1, 2, 3] * 90 , No rotation, 90 counterclockwise, 180 counterclockwise, 270 counterclockwise
            for i in range(4):
                change_image = cv2.warpAffine(result_image, cv2.getRotationMatrix2D((256,256), i*90, 1.0), (512, 512))
                cv2.imwrite("Photos/{}/{}/{}_{}_{}.{}".format(output, _class, filename[0], str(count).zfill(2), str(i+1), filename[1]), change_image)
            
            change_image = cv2.flip(result_image, 0)# 0 flipped vertically, 1 flipped horizontally
            cv2.imwrite("Photos/{}/{}/{}_{}_{}.{}".format(output, _class, filename[0], str(count).zfill(2), "5", filename[1]), change_image)
            
            change_image = cv2.flip(result_image, 1)# 0 flipped vertically, 1 flipped horizontally
            cv2.imwrite("Photos/{}/{}/{}_{}_{}.{}".format(output, _class, filename[0], str(count).zfill(2), "6", filename[1]), change_image)


def makedirs(filename):
    try:
        os.makedirs('Photos' + '/{}'.format(filename))
        os.makedirs('Photos/{}'.format(filename) + '/Benign')
        os.makedirs('Photos/{}'.format(filename) + '/InSitu')
        os.makedirs('Photos/{}'.format(filename) + '/Invasive')
        os.makedirs('Photos/{}'.format(filename) + '/Normal')
    except FileExistsError:
        print("FileExistsError")
        os._exit(0)
    

if __name__=="__main__":
    parser = argparse.ArgumentParser(description='Split image')
    parser.add_argument('--input', '-i', type=str, default='Photos/Normalization/train', help='输入图像路径，文件夹')
    parser.add_argument('--size', '-s', type=int, default=1024, required=True, help='图像分割大小，默认1024')
    parser.add_argument('--num', '-n', type=int, default=12, required=True, help='每张大图增强个数，默认12')
    args = parser.parse_args()
    
    output = "Augmentation_Random_{}_{}".format(args.size, args.num)
    makedirs(output)

    for root, dirs, files in os.walk(args.input):
        for file in files:
            if 'tif' in file:
                split_augment_image_random("{}/{}".format(root,file), file, output, args.size, args.num)